/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.fn.harness.state;



import com.bff.gaia.unified.fn.harness.data.UnifiedFnDataGrpcClient;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.StateRequest;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.StateResponse;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnStateGrpc;

import com.bff.gaia.unified.model.pipeline.v1.Endpoints;

import com.bff.gaia.unified.model.pipeline.v1.Endpoints.ApiServiceDescriptor;

import com.bff.gaia.unified.sdk.fn.IdGenerator;

import com.bff.gaia.unified.sdk.fn.stream.OutboundObserverFactory;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.io.grpc.ManagedChannel;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.io.grpc.stub.StreamObserver;

import org.slf4j.Logger;

import org.slf4j.LoggerFactory;



import java.io.IOException;

import java.util.Map;

import java.util.concurrent.CompletableFuture;

import java.util.concurrent.ConcurrentHashMap;

import java.util.concurrent.ConcurrentMap;

import java.util.function.Function;



/**

 * A cache of {@link UnifiedFnStateClient}s which handle Unified Fn State requests using gRPC.

 *

 * <p>TODO: Add the ability to close which cancels any pending and stops any future requests.

 */

public class UnifiedFnStateGrpcClientCache {

  private static final Logger LOG = LoggerFactory.getLogger(UnifiedFnDataGrpcClient.class);



  private final ConcurrentMap<ApiServiceDescriptor, UnifiedFnStateClient> cache;

  private final Function<ApiServiceDescriptor, ManagedChannel> channelFactory;

  private final OutboundObserverFactory outboundObserverFactory;

  private final IdGenerator idGenerator;



  public UnifiedFnStateGrpcClientCache(

      IdGenerator idGenerator,

      Function<Endpoints.ApiServiceDescriptor, ManagedChannel> channelFactory,

      OutboundObserverFactory outboundObserverFactory) {

    this.idGenerator = idGenerator;

    this.channelFactory = channelFactory;

    this.outboundObserverFactory = outboundObserverFactory;

    this.cache = new ConcurrentHashMap<>();

  }



  /**

   * ( Creates or returns an existing {@link UnifiedFnStateClient} depending on whether the passed in

   * {@link ApiServiceDescriptor} currently has a {@link UnifiedFnStateClient} bound to the same

   * channel.

   */

  public UnifiedFnStateClient forApiServiceDescriptor(ApiServiceDescriptor apiServiceDescriptor)

      throws IOException {

    return cache.computeIfAbsent(apiServiceDescriptor, this::createUnifiedFnStateClient);

  }



  private UnifiedFnStateClient createUnifiedFnStateClient(ApiServiceDescriptor apiServiceDescriptor) {

    return new GrpcStateClient(apiServiceDescriptor);

  }



  /** A {@link UnifiedFnStateClient} for a given {@link ApiServiceDescriptor}. */

  private class GrpcStateClient implements UnifiedFnStateClient {

    private final ApiServiceDescriptor apiServiceDescriptor;

    private final ConcurrentMap<String, CompletableFuture<StateResponse>> outstandingRequests;

    private final StreamObserver<StateRequest> outboundObserver;

    private final ManagedChannel channel;

    private volatile RuntimeException closed;



    private GrpcStateClient(ApiServiceDescriptor apiServiceDescriptor) {

      this.apiServiceDescriptor = apiServiceDescriptor;

      this.outstandingRequests = new ConcurrentHashMap<>();

      this.channel = channelFactory.apply(apiServiceDescriptor);

      this.outboundObserver =

          outboundObserverFactory.outboundObserverFor(

              UnifiedFnStateGrpc.newStub(channel)::state, new InboundObserver());

    }



    @Override

    public void handle(

		StateRequest.Builder requestBuilder, CompletableFuture<StateResponse> response) {

      requestBuilder.setId(idGenerator.getId());

      StateRequest request = requestBuilder.build();

      outstandingRequests.put(request.getId(), response);



      // If the server closes, gRPC will throw an error if onNext is called.

      LOG.debug("Sending StateRequest {}", request);

      outboundObserver.onNext(request);

    }



    private synchronized void closeAndCleanUp(RuntimeException cause) {

      if (closed != null) {

        return;

      }

      cache.remove(apiServiceDescriptor);

      closed = cause;



      // Make a copy of the map to make the view of the outstanding requests consistent.

      Map<String, CompletableFuture<StateResponse>> outstandingRequestsCopy =

          new ConcurrentHashMap<>(outstandingRequests);



      if (outstandingRequestsCopy.isEmpty()) {

        outboundObserver.onCompleted();

        return;

      }



      outstandingRequests.clear();

      LOG.error("UnifiedFnState failed, clearing outstanding requests {}", outstandingRequestsCopy);



      for (CompletableFuture<StateResponse> entry : outstandingRequestsCopy.values()) {

        entry.completeExceptionally(cause);

      }

    }



    /**

     * A {@link StreamObserver} which propagates any server side state request responses by

     * completing the outstanding response future.

     *

     * <p>Also propagates server side failures and closes completing any outstanding requests

     * exceptionally.

     */

    private class InboundObserver implements StreamObserver<StateResponse> {

      @Override

      public void onNext(StateResponse value) {

        LOG.debug("Received StateResponse {}", value);

        CompletableFuture<StateResponse> responseFuture = outstandingRequests.remove(value.getId());

        if (responseFuture != null) {

          if (value.getError().isEmpty()) {

            responseFuture.complete(value);

          } else {

            responseFuture.completeExceptionally(new IllegalStateException(value.getError()));

          }

        }

      }



      @Override

      public void onError(Throwable t) {

        closeAndCleanUp(

            t instanceof RuntimeException ? (RuntimeException) t : new RuntimeException(t));

      }



      @Override

      public void onCompleted() {

        closeAndCleanUp(new RuntimeException("Server hanged up."));

      }

    }

  }

}